{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "view-in-github"
   },
   "source": [
    "<a href=\"https://colab.research.google.com/github/dlsyscourse/public_notebooks/blob/main/24_machine_learning_compilation_deployment_implementation.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Mpn1ti5Urdsv"
   },
   "source": [
    "# Lecture 24: Machine Learning Compiler and Deployment\n",
    "\n",
    "In this lecture, we will walk you through some example usage of the machine learning compiler Apache TVM. To learn more, checkout https://tvm.apache.org/\n",
    "\n",
    "The content of this lecture is adapted from TVM's tutorials.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "qXysoqn-vZuF"
   },
   "source": [
    "## Install package\n",
    "\n",
    "To get started, we need to obtain a version of TVM. For quick demo purpose we will use the following command to install a latest version of the TVM unity compiler and related language model dependenchy solution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Xe3vClsD9jlq",
    "outputId": "29482ab0-1e2d-4e99-d1e8-33c8eef1fb4a"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Looking in links: https://mlc.ai/wheels\n",
      "\u001b[33mWARNING: Retrying (Retry(total=4, connect=None, read=None, redirect=None, status=None)) after connection broken by 'ReadTimeoutError(\"HTTPSConnectionPool(host='mlc.ai', port=443): Read timed out. (read timeout=15)\")': /wheels\u001b[0m\u001b[33m\n",
      "\u001b[0m\u001b[33mWARNING: Retrying (Retry(total=3, connect=None, read=None, redirect=None, status=None)) after connection broken by 'ReadTimeoutError(\"HTTPSConnectionPool(host='mlc.ai', port=443): Read timed out. (read timeout=15)\")': /wheels\u001b[0m\u001b[33m\n",
      "\u001b[0m\u001b[33mWARNING: Retrying (Retry(total=2, connect=None, read=None, redirect=None, status=None)) after connection broken by 'ReadTimeoutError(\"HTTPSConnectionPool(host='mlc.ai', port=443): Read timed out. (read timeout=15)\")': /wheels\u001b[0m\u001b[33m\n",
      "\u001b[0m\u001b[33mWARNING: Retrying (Retry(total=1, connect=None, read=None, redirect=None, status=None)) after connection broken by 'ReadTimeoutError(\"HTTPSConnectionPool(host='mlc.ai', port=443): Read timed out. (read timeout=15)\")': /wheels\u001b[0m\u001b[33m\n",
      "\u001b[0m\u001b[33mWARNING: Retrying (Retry(total=0, connect=None, read=None, redirect=None, status=None)) after connection broken by 'ReadTimeoutError(\"HTTPSConnectionPool(host='mlc.ai', port=443): Read timed out. (read timeout=15)\")': /wheels\u001b[0m\u001b[33m\n",
      "\u001b[0m\u001b[33mWARNING: Retrying (Retry(total=4, connect=None, read=None, redirect=None, status=None)) after connection broken by 'ReadTimeoutError(\"HTTPSConnectionPool(host='pypi.org', port=443): Read timed out. (read timeout=15)\")': /simple/mlc-ai-nightly-cu118/\u001b[0m\u001b[33m\n",
      "\u001b[0m\u001b[33mWARNING: Retrying (Retry(total=3, connect=None, read=None, redirect=None, status=None)) after connection broken by 'ReadTimeoutError(\"HTTPSConnectionPool(host='pypi.org', port=443): Read timed out. (read timeout=15)\")': /simple/mlc-ai-nightly-cu118/\u001b[0m\u001b[33m\n",
      "\u001b[0m\u001b[33mWARNING: Retrying (Retry(total=2, connect=None, read=None, redirect=None, status=None)) after connection broken by 'ReadTimeoutError(\"HTTPSConnectionPool(host='pypi.org', port=443): Read timed out. (read timeout=15)\")': /simple/mlc-ai-nightly-cu118/\u001b[0m\u001b[33m\n",
      "\u001b[0m\u001b[33mWARNING: Retrying (Retry(total=1, connect=None, read=None, redirect=None, status=None)) after connection broken by 'ReadTimeoutError(\"HTTPSConnectionPool(host='pypi.org', port=443): Read timed out. (read timeout=15)\")': /simple/mlc-ai-nightly-cu118/\u001b[0m\u001b[33m\n",
      "\u001b[0m\u001b[33mWARNING: Retrying (Retry(total=0, connect=None, read=None, redirect=None, status=None)) after connection broken by 'ReadTimeoutError(\"HTTPSConnectionPool(host='pypi.org', port=443): Read timed out. (read timeout=15)\")': /simple/mlc-ai-nightly-cu118/\u001b[0m\u001b[33m\n",
      "\u001b[0m\u001b[31mERROR: Could not find a version that satisfies the requirement mlc-ai-nightly-cu118 (from versions: none)\u001b[0m\u001b[31m\n",
      "\u001b[0m\u001b[31mERROR: No matching distribution found for mlc-ai-nightly-cu118\u001b[0m\u001b[31m\n",
      "\u001b[0m"
     ]
    }
   ],
   "source": [
    "!pip install --pre  mlc-ai-nightly-cu118 mlc-chat-nightly-cu118 -f https://mlc.ai/wheels"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "BBIuE2jc1DaU"
   },
   "source": [
    "## Loop-level representation and transformations\n",
    "\n",
    "Let us start with a vector add example. the follow code snippet allows us to create a vector add code, and store it in a container called IRModule."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "vvfOgcu-YdaB"
   },
   "outputs": [],
   "source": [
    "import tvm\n",
    "from tvm.ir.module import IRModule\n",
    "from tvm.script import tir as T\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "xxeOOes6OEkw"
   },
   "outputs": [],
   "source": [
    "def lnumpy_add(a, b, c):\n",
    "    for i in range(128):\n",
    "        c[i] = a[i] + b[i]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "qCViJNUNYfTW"
   },
   "outputs": [],
   "source": [
    "from tvm import te\n",
    "\n",
    "A = te.placeholder(shape=(128,), dtype=\"float32\", name=\"A\")\n",
    "B = te.placeholder(shape=(128,), dtype=\"float32\", name=\"B\")\n",
    "C = te.compute((128,), lambda i: A[i] + B[i], name=\"C\")\n",
    "func = te.create_prim_func([A, B, C])\n",
    "func = func.with_attr(\"global_symbol\", \"main\")\n",
    "ir_module = IRModule({\"main\": func})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4PJd0Pw8zVQD"
   },
   "source": [
    "An IRModule contains a collection of low-level functions, we can use the script function to inspect the functions inside an IRModule.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 280
    },
    "id": "VXy-4v3Czax9",
    "outputId": "1192722c-f435-4fb1-9c4f-8ce25cef9731"
   },
   "outputs": [],
   "source": [
    "ir_module.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "tOMSOuW6aIJg"
   },
   "source": [
    "### Build and run\n",
    "\n",
    "We can turn the programs in an IRModule to runnable functions by calling a build function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "oESoqN-xaTCf",
    "outputId": "245a8e09-51f8-46b2-b22c-c122753ca8d2"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'tvm.driver.build_module.OperatorModule'>\n"
     ]
    }
   ],
   "source": [
    "rt_mod = tvm.build(ir_module, target=\"llvm\")  # The module for CPU backends.\n",
    "print(type(rt_mod))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Y2ZfGrH1z6SV"
   },
   "source": [
    "After build, mod contains a collection of runnable functions. We can retrieve each function by its name."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "id": "5I3GqwnRz-Ne"
   },
   "outputs": [],
   "source": [
    "func = rt_mod[\"main\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "bngdW1eVl683"
   },
   "outputs": [],
   "source": [
    "func"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "id": "DKxo8uq_mNlp"
   },
   "outputs": [],
   "source": [
    "a = tvm.nd.array(np.arange(128, dtype=\"float32\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "id": "1hAFAqv_mP8W"
   },
   "outputs": [],
   "source": [
    "b = tvm.nd.array(np.ones(128, dtype=\"float32\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "id": "TseB1UBumivT"
   },
   "outputs": [],
   "source": [
    "c = tvm.nd.empty((128,), dtype=\"float32\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "heOAWvPAm57P",
    "outputId": "86192aff-f415-4ddb-e9b2-7769f8d226e8"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tvm.nd.NDArray shape=(128,), cpu(0)>\n",
       "array([ 1.5414283e-44,  0.0000000e+00, -1.8209923e+38,  5.9951127e-36,\n",
       "        9.3680810e-38,            nan, -1.7944576e+38,            nan,\n",
       "       -1.7950287e+38,  1.5046100e-36, -2.3580977e-13,  5.0714873e+03,\n",
       "       -2.1031324e-29,  4.3239867e-41, -1.0050381e-27,  4.3239867e-41,\n",
       "       -5.0817178e-28,  7.4720040e-31, -5.4709796e-29,  4.3239867e-41,\n",
       "       -3.8291387e-29,  4.3239867e-41, -3.0956019e+08,  3.5900357e+37,\n",
       "       -2.1032094e-29,  4.3239867e-41, -2.1031901e-29,  4.3239867e-41,\n",
       "        1.0217351e+24,  2.9740162e-25, -2.9161776e-28,  4.3239867e-41,\n",
       "       -1.0050535e-27,  4.3239867e-41,  5.3499954e+35,  2.2136460e-08,\n",
       "       -7.9329131e-28,  4.3239867e-41,  2.1864253e-07,  3.2584393e-41,\n",
       "        3.7308126e+18, -1.3384625e-12, -2.1032190e-29,  4.3239867e-41,\n",
       "       -1.0052569e-27,  4.3239867e-41,  2.3641396e+20,  1.1002859e+12,\n",
       "       -2.1032286e-29,  4.3239867e-41, -1.0522280e-27,  4.3239867e-41,\n",
       "       -1.2803440e-37,  7.2564893e-31, -2.1032383e-29,  4.3239867e-41,\n",
       "       -1.0522156e-27,  4.3239867e-41,  2.2818005e-05,  4.0465242e-08,\n",
       "       -5.7715267e-28,  4.3239867e-41, -9.1011591e-28,  4.3239867e-41,\n",
       "        1.5972619e-25,  1.1333209e+00, -7.3430239e-28,  4.3239867e-41,\n",
       "        2.1928236e-07,  3.2584393e-41,  1.3122041e+07,  2.4848557e+26,\n",
       "       -1.9644851e-28,  4.3239867e-41, -1.0074093e-27,  4.3239867e-41,\n",
       "        0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,\n",
       "        0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,\n",
       "        0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,\n",
       "        0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,\n",
       "        0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,\n",
       "        0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,\n",
       "        0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,\n",
       "        0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,\n",
       "        0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,\n",
       "        0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,\n",
       "        0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,\n",
       "        0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,\n",
       "        0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00],\n",
       "      dtype=float32)"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    " c"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "p68xZ0_P0MPw"
   },
   "source": [
    "To invoke the function, we can create three NDArrays in the tvm runtime, and then invoke the generated function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "id": "SMkcgO-L0Xr5"
   },
   "outputs": [],
   "source": [
    "func(a, b, c)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "AakkTpE50b6o",
    "outputId": "db4a5861-d341-4672-d532-a72e34cc0ed1"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[  0.   1.   2.   3.   4.   5.   6.   7.   8.   9.  10.  11.  12.  13.\n",
      "  14.  15.  16.  17.  18.  19.  20.  21.  22.  23.  24.  25.  26.  27.\n",
      "  28.  29.  30.  31.  32.  33.  34.  35.  36.  37.  38.  39.  40.  41.\n",
      "  42.  43.  44.  45.  46.  47.  48.  49.  50.  51.  52.  53.  54.  55.\n",
      "  56.  57.  58.  59.  60.  61.  62.  63.  64.  65.  66.  67.  68.  69.\n",
      "  70.  71.  72.  73.  74.  75.  76.  77.  78.  79.  80.  81.  82.  83.\n",
      "  84.  85.  86.  87.  88.  89.  90.  91.  92.  93.  94.  95.  96.  97.\n",
      "  98.  99. 100. 101. 102. 103. 104. 105. 106. 107. 108. 109. 110. 111.\n",
      " 112. 113. 114. 115. 116. 117. 118. 119. 120. 121. 122. 123. 124. 125.\n",
      " 126. 127.]\n",
      "[1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1.\n",
      " 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1.\n",
      " 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1.\n",
      " 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1.\n",
      " 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1.\n",
      " 1. 1. 1. 1. 1. 1. 1. 1.]\n",
      "[  1.   2.   3.   4.   5.   6.   7.   8.   9.  10.  11.  12.  13.  14.\n",
      "  15.  16.  17.  18.  19.  20.  21.  22.  23.  24.  25.  26.  27.  28.\n",
      "  29.  30.  31.  32.  33.  34.  35.  36.  37.  38.  39.  40.  41.  42.\n",
      "  43.  44.  45.  46.  47.  48.  49.  50.  51.  52.  53.  54.  55.  56.\n",
      "  57.  58.  59.  60.  61.  62.  63.  64.  65.  66.  67.  68.  69.  70.\n",
      "  71.  72.  73.  74.  75.  76.  77.  78.  79.  80.  81.  82.  83.  84.\n",
      "  85.  86.  87.  88.  89.  90.  91.  92.  93.  94.  95.  96.  97.  98.\n",
      "  99. 100. 101. 102. 103. 104. 105. 106. 107. 108. 109. 110. 111. 112.\n",
      " 113. 114. 115. 116. 117. 118. 119. 120. 121. 122. 123. 124. 125. 126.\n",
      " 127. 128.]\n"
     ]
    }
   ],
   "source": [
    "print(a)\n",
    "print(b)\n",
    "print(c)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "i_MIDZCOcmwp"
   },
   "source": [
    "### Transform the code\n",
    "\n",
    "The IRModule is the central data structure for program optimization, which can be transformed by a helper class called Schedule. A schedule contains multiple primitive methods to interactively transform the program. Each primitive transforms the program in certain ways to bring additional performance optimizations."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "qlc8oWsZc0_p"
   },
   "source": [
    "Let us try to transform the module, we can do it by creating a Schedule instance.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "id": "GEwPVl8AOa7i"
   },
   "outputs": [],
   "source": [
    "sch = tvm.tir.Schedule(ir_module)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 280
    },
    "id": "-6DtngEJShsR",
    "outputId": "92f85b5f-a2bb-47b1-cc39-81c0889968f0"
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
       "\n",
       "<span style=\"color: #AA22FF\">@I</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>ir_module\n",
       "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #0000FF; font-weight: bold\">Module</span>:\n",
       "    <span style=\"color: #AA22FF\">@T</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>prim_func\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">main</span>(A: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">128</span>,), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), B: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">128</span>,), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), C: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">128</span>,), <span style=\"color: #BA2121\">&quot;float32&quot;</span>)):\n",
       "        T<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>bool(<span style=\"color: #008000; font-weight: bold\">True</span>)})\n",
       "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
       "        <span style=\"color: #008000; font-weight: bold\">for</span> i <span style=\"color: #008000; font-weight: bold\">in</span> range(<span style=\"color: #008000\">128</span>):\n",
       "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;C&quot;</span>):\n",
       "                v_i <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>spatial(<span style=\"color: #008000\">128</span>, i)\n",
       "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads(A[v_i], B[v_i])\n",
       "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(C[v_i])\n",
       "                C[v_i] <span style=\"color: #AA22FF; font-weight: bold\">=</span> A[v_i] <span style=\"color: #AA22FF; font-weight: bold\">+</span> B[v_i]\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "sch.mod.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {
    "id": "48-aWNqxSeMW"
   },
   "outputs": [],
   "source": [
    "blockC = sch.get_block(\"C\")\n",
    "i, = sch.get_loops(blockC)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {
    "id": "d-fxPJCHSeQQ"
   },
   "outputs": [],
   "source": [
    "i0, i1 = sch.split(i, factors=[None, 8])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 280
    },
    "id": "jXQKYSfQOdiN",
    "outputId": "25a87b6b-d59b-4741-dc81-1c7d72003d1e"
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
       "\n",
       "<span style=\"color: #AA22FF\">@I</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>ir_module\n",
       "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #0000FF; font-weight: bold\">Module</span>:\n",
       "    <span style=\"color: #AA22FF\">@T</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>prim_func\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">main</span>(A: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">128</span>,), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), B: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">128</span>,), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), C: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">128</span>,), <span style=\"color: #BA2121\">&quot;float32&quot;</span>)):\n",
       "        T<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>bool(<span style=\"color: #008000; font-weight: bold\">True</span>)})\n",
       "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
       "        <span style=\"color: #008000; font-weight: bold\">for</span> i_0, i_1 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>grid(<span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">8</span>):\n",
       "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;C&quot;</span>):\n",
       "                v_i <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>spatial(<span style=\"color: #008000\">128</span>, i_0 <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">8</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> i_1)\n",
       "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads(A[v_i], B[v_i])\n",
       "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(C[v_i])\n",
       "                C[v_i] <span style=\"color: #AA22FF; font-weight: bold\">=</span> A[v_i] <span style=\"color: #AA22FF; font-weight: bold\">+</span> B[v_i]\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "sch.mod.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {
    "id": "l5EySwKYObld"
   },
   "outputs": [],
   "source": [
    "sch.reorder(i1, i0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 280
    },
    "id": "lfZMmgJ3Obr7",
    "outputId": "8a4a4bae-bee7-4a10-c70f-7955ca92c629"
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
       "\n",
       "<span style=\"color: #AA22FF\">@I</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>ir_module\n",
       "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #0000FF; font-weight: bold\">Module</span>:\n",
       "    <span style=\"color: #AA22FF\">@T</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>prim_func\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">main</span>(A: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">128</span>,), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), B: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">128</span>,), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), C: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">128</span>,), <span style=\"color: #BA2121\">&quot;float32&quot;</span>)):\n",
       "        T<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>bool(<span style=\"color: #008000; font-weight: bold\">True</span>)})\n",
       "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
       "        <span style=\"color: #008000; font-weight: bold\">for</span> i_1, i_0 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>grid(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">16</span>):\n",
       "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;C&quot;</span>):\n",
       "                v_i <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>spatial(<span style=\"color: #008000\">128</span>, i_0 <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">8</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> i_1)\n",
       "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads(A[v_i], B[v_i])\n",
       "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(C[v_i])\n",
       "                C[v_i] <span style=\"color: #AA22FF; font-weight: bold\">=</span> A[v_i] <span style=\"color: #AA22FF; font-weight: bold\">+</span> B[v_i]\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "sch.mod.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {
    "id": "H8n1BTXBObvD"
   },
   "outputs": [],
   "source": [
    "sch.reorder(i0, i1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 280
    },
    "id": "-StL-Jx0ObyK",
    "outputId": "5f9956bf-e063-45dd-e1b6-f3b90af34429"
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
       "\n",
       "<span style=\"color: #AA22FF\">@I</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>ir_module\n",
       "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #0000FF; font-weight: bold\">Module</span>:\n",
       "    <span style=\"color: #AA22FF\">@T</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>prim_func\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">main</span>(A: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">128</span>,), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), B: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">128</span>,), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), C: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">128</span>,), <span style=\"color: #BA2121\">&quot;float32&quot;</span>)):\n",
       "        T<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>bool(<span style=\"color: #008000; font-weight: bold\">True</span>)})\n",
       "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
       "        <span style=\"color: #008000; font-weight: bold\">for</span> i_0, i_1 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>grid(<span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">8</span>):\n",
       "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;C&quot;</span>):\n",
       "                v_i <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>spatial(<span style=\"color: #008000\">128</span>, i_0 <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">8</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> i_1)\n",
       "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads(A[v_i], B[v_i])\n",
       "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(C[v_i])\n",
       "                C[v_i] <span style=\"color: #AA22FF; font-weight: bold\">=</span> A[v_i] <span style=\"color: #AA22FF; font-weight: bold\">+</span> B[v_i]\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "sch.mod.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {
    "id": "rGv-EfTuTbXM"
   },
   "outputs": [],
   "source": [
    "sch.parallel(i0)\n",
    "sch.vectorize(i1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 297
    },
    "id": "wl2uOxAATbZ6",
    "outputId": "89fe1278-e0f9-4e87-85e9-bdfa1aa976cf"
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
       "\n",
       "<span style=\"color: #AA22FF\">@I</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>ir_module\n",
       "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #0000FF; font-weight: bold\">Module</span>:\n",
       "    <span style=\"color: #AA22FF\">@T</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>prim_func\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">main</span>(A: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">128</span>,), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), B: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">128</span>,), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), C: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">128</span>,), <span style=\"color: #BA2121\">&quot;float32&quot;</span>)):\n",
       "        T<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>bool(<span style=\"color: #008000; font-weight: bold\">True</span>)})\n",
       "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
       "        <span style=\"color: #008000; font-weight: bold\">for</span> i_0 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>parallel(<span style=\"color: #008000\">16</span>):\n",
       "            <span style=\"color: #008000; font-weight: bold\">for</span> i_1 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>vectorized(<span style=\"color: #008000\">8</span>):\n",
       "                <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;C&quot;</span>):\n",
       "                    v_i <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>spatial(<span style=\"color: #008000\">128</span>, i_0 <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">8</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> i_1)\n",
       "                    T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads(A[v_i], B[v_i])\n",
       "                    T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(C[v_i])\n",
       "                    C[v_i] <span style=\"color: #AA22FF; font-weight: bold\">=</span> A[v_i] <span style=\"color: #AA22FF; font-weight: bold\">+</span> B[v_i]\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "sch.mod.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "V0hCoH-0Tbcq"
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "KuJ3GNa7Tbe9"
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "xwyjwh51cjWI",
    "outputId": "8f25bd05-b6d7-4c41-bd81-b91ec2d0f8c6"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'tvm.tir.schedule.schedule.Schedule'>\n"
     ]
    }
   ],
   "source": [
    "sch = tvm.tir.Schedule(ir_module)\n",
    "print(type(sch))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 280
    },
    "id": "wpiFYFSNVIWa",
    "outputId": "bac1fd55-6358-491e-da8b-9c8cb9e3d9d3"
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
       "\n",
       "<span style=\"color: #AA22FF\">@I</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>ir_module\n",
       "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #0000FF; font-weight: bold\">Module</span>:\n",
       "    <span style=\"color: #AA22FF\">@T</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>prim_func\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">main</span>(A: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">128</span>,), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), B: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">128</span>,), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), C: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">128</span>,), <span style=\"color: #BA2121\">&quot;float32&quot;</span>)):\n",
       "        T<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>bool(<span style=\"color: #008000; font-weight: bold\">True</span>)})\n",
       "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
       "        <span style=\"color: #008000; font-weight: bold\">for</span> i <span style=\"color: #008000; font-weight: bold\">in</span> range(<span style=\"color: #008000\">128</span>):\n",
       "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;C&quot;</span>):\n",
       "                v_i <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>spatial(<span style=\"color: #008000\">128</span>, i)\n",
       "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads(A[v_i], B[v_i])\n",
       "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(C[v_i])\n",
       "                C[v_i] <span style=\"color: #AA22FF; font-weight: bold\">=</span> A[v_i] <span style=\"color: #AA22FF; font-weight: bold\">+</span> B[v_i]\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "ir_module.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "KbtzPuFCU6yR"
   },
   "outputs": [],
   "source": [
    "sch = tvm.tir.Schedule(ir_module)\n",
    "block_C = sch.get_block(\"C\")\n",
    "i, = sch.get_loops(block_C)\n",
    "i0, i1 = sch.split(i, [None, 8])\n",
    "sch.reorder(i1, i0)\n",
    "sch.parallel(i1)\n",
    "sch.mod.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "jUgTV7aZU83o"
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Dw7Fgw8o8HPm"
   },
   "source": [
    "Let us first try to split the loops"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 280
    },
    "id": "kNQf8D0ic4me",
    "outputId": "866c1e70-25c4-47cf-c40d-49d324d921a1"
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
       "\n",
       "<span style=\"color: #AA22FF\">@I</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>ir_module\n",
       "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #0000FF; font-weight: bold\">Module</span>:\n",
       "    <span style=\"color: #AA22FF\">@T</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>prim_func\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">main</span>(A: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">128</span>,), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), B: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">128</span>,), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), C: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">128</span>,), <span style=\"color: #BA2121\">&quot;float32&quot;</span>)):\n",
       "        T<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>bool(<span style=\"color: #008000; font-weight: bold\">True</span>)})\n",
       "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
       "        <span style=\"color: #008000; font-weight: bold\">for</span> i_0, i_1, i_2 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>grid(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">4</span>, <span style=\"color: #008000\">4</span>):\n",
       "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;C&quot;</span>):\n",
       "                v_i <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>spatial(<span style=\"color: #008000\">128</span>, i_0 <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">16</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> i_1 <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">4</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> i_2)\n",
       "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads(A[v_i], B[v_i])\n",
       "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(C[v_i])\n",
       "                C[v_i] <span style=\"color: #AA22FF; font-weight: bold\">=</span> A[v_i] <span style=\"color: #AA22FF; font-weight: bold\">+</span> B[v_i]\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Get block by its name\n",
    "block_c = sch.get_block(\"C\")\n",
    "# Get loops surronding the block\n",
    "(i,) = sch.get_loops(block_c)\n",
    "# Tile the loop nesting.\n",
    "i_0, i_1, i_2 = sch.split(i, factors=[None, 4, 4])\n",
    "sch.mod.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "nzrbvqBSdC-D"
   },
   "source": [
    "We can also reorder the loops, swapping the order of i_0 and i_1\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 280
    },
    "id": "yJWBq7lRdDmn",
    "outputId": "5253c3fe-1b57-48c3-ce71-ec983e85a537"
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
       "\n",
       "<span style=\"color: #AA22FF\">@I</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>ir_module\n",
       "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #0000FF; font-weight: bold\">Module</span>:\n",
       "    <span style=\"color: #AA22FF\">@T</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>prim_func\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">main</span>(A: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">128</span>,), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), B: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">128</span>,), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), C: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">128</span>,), <span style=\"color: #BA2121\">&quot;float32&quot;</span>)):\n",
       "        T<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>bool(<span style=\"color: #008000; font-weight: bold\">True</span>)})\n",
       "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
       "        <span style=\"color: #008000; font-weight: bold\">for</span> i_1, i_0, i_2 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>grid(<span style=\"color: #008000\">4</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">4</span>):\n",
       "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;C&quot;</span>):\n",
       "                v_i <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>spatial(<span style=\"color: #008000\">128</span>, i_0 <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">16</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> i_1 <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">4</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> i_2)\n",
       "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads(A[v_i], B[v_i])\n",
       "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(C[v_i])\n",
       "                C[v_i] <span style=\"color: #AA22FF; font-weight: bold\">=</span> A[v_i] <span style=\"color: #AA22FF; font-weight: bold\">+</span> B[v_i]\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "sch.reorder(i_1, i_0, i_2)\n",
    "sch.mod.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "UmUr6b_L07-8"
   },
   "source": [
    "Finally, we can add hints to the program generator that we want to vectorize the inner most loop."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 297
    },
    "id": "u95zQFuldHs_",
    "outputId": "281a259e-5138-4f90-9793-1a1cdc6472b4"
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
       "\n",
       "<span style=\"color: #AA22FF\">@I</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>ir_module\n",
       "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #0000FF; font-weight: bold\">Module</span>:\n",
       "    <span style=\"color: #AA22FF\">@T</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>prim_func\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">main</span>(A: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">128</span>,), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), B: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">128</span>,), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), C: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">128</span>,), <span style=\"color: #BA2121\">&quot;float32&quot;</span>)):\n",
       "        T<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>bool(<span style=\"color: #008000; font-weight: bold\">True</span>)})\n",
       "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
       "        <span style=\"color: #008000; font-weight: bold\">for</span> i_1, i_0 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>grid(<span style=\"color: #008000\">4</span>, <span style=\"color: #008000\">8</span>):\n",
       "            <span style=\"color: #008000; font-weight: bold\">for</span> i_2 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>vectorized(<span style=\"color: #008000\">4</span>):\n",
       "                <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;C&quot;</span>):\n",
       "                    v_i <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>spatial(<span style=\"color: #008000\">128</span>, i_0 <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">16</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> i_1 <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">4</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> i_2)\n",
       "                    T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads(A[v_i], B[v_i])\n",
       "                    T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(C[v_i])\n",
       "                    C[v_i] <span style=\"color: #AA22FF; font-weight: bold\">=</span> A[v_i] <span style=\"color: #AA22FF; font-weight: bold\">+</span> B[v_i]\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "sch.vectorize(i_2)\n",
    "sch.mod.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "GEqpO14Lf0Lq"
   },
   "source": [
    "## Transforming a matrix multiplication program\n",
    "\n",
    "In the above example, we showed how to transform an vector add. Now let us try to apply that to a slightly more complicated program(matrix multiplication).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 315
    },
    "id": "B3a2bl2GWwMZ",
    "outputId": "7dcea30c-c00f-4944-8814-f5e4c3c4e0ce"
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
       "\n",
       "<span style=\"color: #AA22FF\">@I</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>ir_module\n",
       "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #0000FF; font-weight: bold\">Module</span>:\n",
       "    <span style=\"color: #AA22FF\">@T</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>prim_func\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">main</span>(A: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">1024</span>, <span style=\"color: #008000\">1024</span>), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), B: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">1024</span>, <span style=\"color: #008000\">1024</span>), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), C: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">1024</span>, <span style=\"color: #008000\">1024</span>), <span style=\"color: #BA2121\">&quot;float32&quot;</span>)):\n",
       "        T<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>bool(<span style=\"color: #008000; font-weight: bold\">True</span>)})\n",
       "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
       "        <span style=\"color: #008000; font-weight: bold\">for</span> m, n, k <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>grid(<span style=\"color: #008000\">1024</span>, <span style=\"color: #008000\">1024</span>, <span style=\"color: #008000\">1024</span>):\n",
       "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;C&quot;</span>):\n",
       "                v_m, v_n, v_k <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>remap(<span style=\"color: #BA2121\">&quot;SSR&quot;</span>, [m, n, k])\n",
       "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads(A[v_m, v_k], B[v_k, v_n])\n",
       "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(C[v_m, v_n])\n",
       "                <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>init():\n",
       "                    C[v_m, v_n] <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>float32(<span style=\"color: #008000\">0</span>)\n",
       "                C[v_m, v_n] <span style=\"color: #AA22FF; font-weight: bold\">=</span> C[v_m, v_n] <span style=\"color: #AA22FF; font-weight: bold\">+</span> A[v_m, v_k] <span style=\"color: #AA22FF; font-weight: bold\">*</span> B[v_k, v_n]\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "M = 1024\n",
    "K = 1024\n",
    "N = 1024\n",
    "\n",
    "# The default tensor type in tvm\n",
    "dtype = \"float32\"\n",
    "\n",
    "target = \"llvm\"\n",
    "dev = tvm.device(target, 0)\n",
    "\n",
    "# Algorithm\n",
    "k = te.reduce_axis((0, K), \"k\")\n",
    "A = te.placeholder((M, K), name=\"A\")\n",
    "B = te.placeholder((K, N), name=\"B\")\n",
    "C = te.compute((M, N), lambda m, n: te.sum(A[m, k] * B[k, n], axis=k), name=\"C\")\n",
    "\n",
    "# Default schedule\n",
    "func = te.create_prim_func([A, B, C])\n",
    "func = func.with_attr(\"global_symbol\", \"main\")\n",
    "ir_module = IRModule({\"main\": func})\n",
    "ir_module.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "iGaL3GMMXWNX",
    "outputId": "cc1883cc-cad1-44c7-c2e5-70d88fc0ba11"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Baseline time cost 3.08041 sec\n"
     ]
    }
   ],
   "source": [
    "func = tvm.build(ir_module, target=\"llvm\")  # The module for CPU backends.\n",
    "a_np = np.random.rand(M, K).astype(dtype)\n",
    "b_np = np.random.rand(K, N).astype(dtype)\n",
    "a = tvm.nd.array(a_np, dev)\n",
    "b = tvm.nd.array(b_np, dev)\n",
    "c = tvm.nd.array(np.zeros((M, N), dtype=dtype), dev)\n",
    "func(a, b, c)\n",
    "\n",
    "evaluator = func.time_evaluator(\"main\", dev, number=3)\n",
    "print(\"Baseline time cost %g sec\" % evaluator(a, b, c).mean)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "pSYr3MJ8Vkzf",
    "outputId": "f86b2872-2a93-4598-e6a3-cce7eecdcc5a"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[ 0.0000000e+00, -1.2207031e-04,  0.0000000e+00, ...,\n",
       "         3.0517578e-05,  0.0000000e+00,  1.2207031e-04],\n",
       "       [ 6.1035156e-05, -9.1552734e-05,  3.0517578e-05, ...,\n",
       "        -9.1552734e-05, -1.5258789e-05,  0.0000000e+00],\n",
       "       [-3.0517578e-05,  1.5258789e-04,  3.0517578e-05, ...,\n",
       "         6.1035156e-05, -1.5258789e-04,  0.0000000e+00],\n",
       "       ...,\n",
       "       [-6.1035156e-05, -1.2207031e-04,  9.1552734e-05, ...,\n",
       "        -1.2207031e-04, -1.3732910e-04,  1.2207031e-04],\n",
       "       [-9.1552734e-05,  0.0000000e+00,  6.1035156e-05, ...,\n",
       "         3.0517578e-05, -1.6784668e-04, -3.0517578e-05],\n",
       "       [ 2.4414062e-04, -1.5258789e-04,  1.5258789e-04, ...,\n",
       "        -1.3732910e-04, -1.2207031e-04, -7.6293945e-05]], dtype=float32)"
      ]
     },
     "execution_count": 48,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "c.numpy() - a_np @b_np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 315
    },
    "id": "PsTxdgEWXWie",
    "outputId": "58b85d3d-dcbb-408f-cede-a111c58fe435"
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
       "\n",
       "<span style=\"color: #AA22FF\">@I</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>ir_module\n",
       "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #0000FF; font-weight: bold\">Module</span>:\n",
       "    <span style=\"color: #AA22FF\">@T</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>prim_func\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">main</span>(A: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">1024</span>, <span style=\"color: #008000\">1024</span>), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), B: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">1024</span>, <span style=\"color: #008000\">1024</span>), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), C: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">1024</span>, <span style=\"color: #008000\">1024</span>), <span style=\"color: #BA2121\">&quot;float32&quot;</span>)):\n",
       "        T<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>bool(<span style=\"color: #008000; font-weight: bold\">True</span>)})\n",
       "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
       "        <span style=\"color: #008000; font-weight: bold\">for</span> m, n, k <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>grid(<span style=\"color: #008000\">1024</span>, <span style=\"color: #008000\">1024</span>, <span style=\"color: #008000\">1024</span>):\n",
       "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;C&quot;</span>):\n",
       "                v_m, v_n, v_k <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>remap(<span style=\"color: #BA2121\">&quot;SSR&quot;</span>, [m, n, k])\n",
       "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads(A[v_m, v_k], B[v_k, v_n])\n",
       "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(C[v_m, v_n])\n",
       "                <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>init():\n",
       "                    C[v_m, v_n] <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>float32(<span style=\"color: #008000\">0</span>)\n",
       "                C[v_m, v_n] <span style=\"color: #AA22FF; font-weight: bold\">=</span> C[v_m, v_n] <span style=\"color: #AA22FF; font-weight: bold\">+</span> A[v_m, v_k] <span style=\"color: #AA22FF; font-weight: bold\">*</span> B[v_k, v_n]\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "ir_module.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "swj-gMz-1vBE"
   },
   "source": [
    "We can transform the loop access pattern to make it more cache friendly. Let us use the following schedule."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 350
    },
    "id": "erNEAsU-Ofo0",
    "outputId": "699fff51-ed2c-4c14-cc8f-92995c7ae536"
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
       "\n",
       "<span style=\"color: #AA22FF\">@I</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>ir_module\n",
       "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #0000FF; font-weight: bold\">Module</span>:\n",
       "    <span style=\"color: #AA22FF\">@T</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>prim_func\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">main</span>(A: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">1024</span>, <span style=\"color: #008000\">1024</span>), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), B: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">1024</span>, <span style=\"color: #008000\">1024</span>), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), C: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">1024</span>, <span style=\"color: #008000\">1024</span>), <span style=\"color: #BA2121\">&quot;float32&quot;</span>)):\n",
       "        T<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>bool(<span style=\"color: #008000; font-weight: bold\">True</span>)})\n",
       "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
       "        <span style=\"color: #008000; font-weight: bold\">for</span> m_0, n_0, k, m_1, n_1 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>grid(<span style=\"color: #008000\">128</span>, <span style=\"color: #008000\">128</span>, <span style=\"color: #008000\">1024</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>):\n",
       "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;C&quot;</span>):\n",
       "                v_m <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>spatial(<span style=\"color: #008000\">1024</span>, m_0 <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">8</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> m_1)\n",
       "                v_n <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>spatial(<span style=\"color: #008000\">1024</span>, n_0 <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">8</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> n_1)\n",
       "                v_k <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>reduce(<span style=\"color: #008000\">1024</span>, k)\n",
       "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads(A[v_m, v_k], B[v_k, v_n])\n",
       "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(C[v_m, v_n])\n",
       "                <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>init():\n",
       "                    C[v_m, v_n] <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>float32(<span style=\"color: #008000\">0</span>)\n",
       "                C[v_m, v_n] <span style=\"color: #AA22FF; font-weight: bold\">=</span> C[v_m, v_n] <span style=\"color: #AA22FF; font-weight: bold\">+</span> A[v_m, v_k] <span style=\"color: #AA22FF; font-weight: bold\">*</span> B[v_k, v_n]\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "def transform(sch, tile_m, tile_n):\n",
    "    block_C = sch.get_block(\"C\")\n",
    "    m, n, k = sch.get_loops(block_C)\n",
    "    mo, mi = sch.split(m, [None, tile_m])\n",
    "    no, ni = sch.split(n, [None, tile_n])\n",
    "    sch.reorder(mo, no, k, mi, ni)\n",
    "    return sch\n",
    "sch = tvm.tir.Schedule(ir_module)\n",
    "sch = transform(sch, 8, 8)\n",
    "sch.mod.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 367
    },
    "id": "QgR9qO49Oham",
    "outputId": "719054f6-21f1-43c2-f7f6-55d84d1786de"
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
       "\n",
       "<span style=\"color: #AA22FF\">@I</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>ir_module\n",
       "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #0000FF; font-weight: bold\">Module</span>:\n",
       "    <span style=\"color: #AA22FF\">@T</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>prim_func\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">main</span>(A: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">1024</span>, <span style=\"color: #008000\">1024</span>), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), B: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">1024</span>, <span style=\"color: #008000\">1024</span>), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), C: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">1024</span>, <span style=\"color: #008000\">1024</span>), <span style=\"color: #BA2121\">&quot;float32&quot;</span>)):\n",
       "        T<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>bool(<span style=\"color: #008000; font-weight: bold\">True</span>)})\n",
       "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
       "        <span style=\"color: #008000; font-weight: bold\">for</span> m_0, n_0, k, m_1, n_1 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>grid(<span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">1024</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>):\n",
       "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;C&quot;</span>):\n",
       "                v_m <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>spatial(<span style=\"color: #008000\">1024</span>, m_0 <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">32</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> m_1)\n",
       "                v_n <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>spatial(<span style=\"color: #008000\">1024</span>, n_0 <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">32</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> n_1)\n",
       "                v_k <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>reduce(<span style=\"color: #008000\">1024</span>, k)\n",
       "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads(A[v_m, v_k], B[v_k, v_n])\n",
       "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(C[v_m, v_n])\n",
       "                <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>init():\n",
       "                    C[v_m, v_n] <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>float32(<span style=\"color: #008000\">0</span>)\n",
       "                C[v_m, v_n] <span style=\"color: #AA22FF; font-weight: bold\">=</span> C[v_m, v_n] <span style=\"color: #AA22FF; font-weight: bold\">+</span> A[v_m, v_k] <span style=\"color: #AA22FF; font-weight: bold\">*</span> B[v_k, v_n]\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Transformed time cost 0.300325 sec\n"
     ]
    }
   ],
   "source": [
    "sch = tvm.tir.Schedule(ir_module)\n",
    "sch = transform(sch, 32, 32)\n",
    "sch.mod.show()\n",
    "mod = tvm.build(sch.mod, target=\"llvm\")\n",
    "new_eval = mod.time_evaluator(\"main\", number=3, dev=tvm.cpu())\n",
    "print(\"Transformed time cost %g sec\" % new_eval(a, b, c).mean)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "W60q68KRgdNL"
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "h1RQGOBjn4w_"
   },
   "source": [
    "Try to change the value of bn to see what performance you can get. In pratice, we will leverage an automated system to search over a set of possible transfromations to find an optimal one."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "sXQnXaPCnxJH"
   },
   "source": [
    "\n",
    "\n",
    "![image.png]()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "XtGozrXPmlz4"
   },
   "source": [
    "There are other optimizations that can be applied here, such as vectorization, parallelization and data layout optimization. Please checkout"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "d-8iHbr8jp04"
   },
   "source": [
    "## End to end model deployment\n",
    "\n",
    "Finally, let us walk through an example flow for an end to end model deployment.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "wrQ4TYeEBMBb",
    "outputId": "5b548356-38a4-440b-dbc9-4a7dd7f9607a"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Git LFS initialized.\n"
     ]
    }
   ],
   "source": [
    "!git lfs install\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "y-pLH_BdBSwz",
    "outputId": "e06d6e56-f821-4f00-814f-78d2883e1686"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Cloning into 'dist/prebuilt/lib'...\n",
      "remote: Enumerating objects: 389, done.\u001b[K\n",
      "remote: Counting objects: 100% (115/115), done.\u001b[K\n",
      "remote: Compressing objects: 100% (38/38), done.\u001b[K\n",
      "remote: Total 389 (delta 94), reused 93 (delta 77), pack-reused 274\u001b[K\n",
      "Receiving objects: 100% (389/389), 126.00 MiB | 14.01 MiB/s, done.\n",
      "Resolving deltas: 100% (279/279), done.\n",
      "Updating files: 100% (100/100), done.\n"
     ]
    }
   ],
   "source": [
    "!mkdir -p dist/prebuilt\n",
    "!git clone https://github.com/mlc-ai/binary-mlc-llm-libs.git dist/prebuilt/lib"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "F_5Xka8fBVb3",
    "outputId": "95608e23-a068-4ead-d77b-f476b30d9c3c"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Cloning into 'mlc-chat-Llama-2-7b-chat-hf-q4f16_1'...\n",
      "remote: Enumerating objects: 129, done.\u001b[K\n",
      "remote: Counting objects: 100% (3/3), done.\u001b[K\n",
      "remote: Compressing objects: 100% (3/3), done.\u001b[K\n",
      "remote: Total 129 (delta 0), reused 0 (delta 0), pack-reused 126\u001b[K\n",
      "Receiving objects: 100% (129/129), 500.53 KiB | 19.25 MiB/s, done.\n",
      "Filtering content: 100% (116/116), 3.53 GiB | 58.34 MiB/s, done.\n"
     ]
    }
   ],
   "source": [
    "!cd dist/prebuilt && git clone https://huggingface.co/mlc-ai/mlc-chat-Llama-2-7b-chat-hf-q4f16_1\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "id": "KV2PMdivBbla"
   },
   "outputs": [],
   "source": [
    "from mlc_chat import ChatModule\n",
    "from mlc_chat.callback import StreamToStdout\n",
    "\n",
    "cm = ChatModule(model=\"Llama-2-7b-chat-hf-q4f16_1\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "4rTjoMR3BcgI",
    "outputId": "12ca253e-0eb3-445d-9df8-4c993bc8f423"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Hello! I'm glad you asked! Python was first released in 1991 by Guido van Rossum. It was initially called \"Python\" because van Rossum was a fan of the British comedy group Monty Python's Flying Circus. The language was created as a hobby project, and it quickly gained popularity among computer programmers due to its simplicity and ease of use. Since its initial release, Python has undergone many updates and improvements, and it has become one of the most popular programming languages in the world. Is there anything else you would like to know?\n"
     ]
    }
   ],
   "source": [
    "output = cm.generate(\n",
    "    prompt=\"When was Python released?\",\n",
    "    progress_callback=StreamToStdout(callback_interval=2),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "RAlmhbBrBenk",
    "outputId": "67e4c0eb-2c32-4841-cca3-38458843205f"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Prompt: tell me about cmu\n",
      "Carnegie Mellon University (CMU) is a private research university located in Pittsburgh, Pennsylvania, United States. It was founded in 1900 as the Carnegie Technical Schools by Andrew Carnegie, and later merged with the Mellon Institute of Industrial Research in 1967. CMU is highly regarded for its academic excellence, innovative research, and strong connections with industry leaders.\n",
      "Here are some key points about Carnegie Mellon University:\n",
      "1. Academics: CMU offers a wide range of undergraduate and graduate degree programs in fields such as engineering, computer science, robotics, artificial intelligence, business, public policy, and the arts. The university is known for its interdisciplinary approach to education, which allows students to explore multiple areas of interest.\n",
      "2. Research: CMU is a leading research university, with a strong focus on interdisciplinary collaboration and innovation. The university has a long history of groundbreaking research in fields such as computer science, robotics, artificial intelligence, and engineering.\n",
      "3. Location: CMU is located in Pittsburgh, Pennsylvania, which offers a unique blend of urban and suburban living. Pittsburgh is known for its vibrant cultural scene, rich history, and affordable cost of living.\n",
      "4. Size: CMU has a relatively small student body, with around 14,000 students overall, including 8,000 undergraduates. This allows for a more personalized and intimate learning experience.\n",
      "5. Campus culture: CMU has a diverse and inclusive campus culture, with a strong focus on community engagement and social responsibility. The university is committed to creating a welcoming and supportive environment for all students, faculty, and staff.\n",
      "6. Athletics: CMU has a strong athletics program, with teams competing in the NCAA Division III and the University Athletic Association (UAA). The university is known for its men's and women's basketball teams, as well as its men's and women's soccer teams.\n",
      "7. Extracurriculars: CMU has a wide range of extracurricular activities and clubs, including student organizations, cultural groups, and community service organizations.\n",
      "8. Career outcomes: CMU has a strong reputation in the job market, with\n"
     ]
    }
   ],
   "source": [
    "prompt = input(\"Prompt: \")\n",
    "output = cm.generate(prompt=prompt, progress_callback=StreamToStdout(callback_interval=2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "B-SQZ2T8Bi3a",
    "outputId": "fd5dce77-93d4-4db1-b3d6-2dfdd8f60799"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "prefill: 506.3 tok/s, decode: 49.1 tok/s\n"
     ]
    }
   ],
   "source": [
    "print(cm.stats())\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "40rcbuOzBk5t"
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "MtJ1iz39BYQn"
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "authorship_tag": "ABX9TyNvCvMpsPaukm6f9gppnhHH",
   "gpuType": "T4",
   "include_colab_link": true,
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
